// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#pragma once

#ifndef _PROTECTED_FS_H_
#define _PROTECTED_FS_H_

#include "protected_fs_config.h"

#include "protected_fs_nodes.h"
#include "lru_cache.h"
#ifdef NON_SGX_PROTECTED_FS
#include "non_sgx_protected_fs.h"
#else
#include "sgx_error.h"
#include "sgx_tcrypto.h"
#include "errno.h"
#include <sgx_thread.h>
#endif
#include "sgx_tprotected_fs.h"

typedef enum
{
	SGX_FILE_STATUS_OK = 0,
	SGX_FILE_STATUS_NOT_INITIALIZED,
	SGX_FILE_STATUS_FLUSH_ERROR,
	SGX_FILE_STATUS_WRITE_TO_DISK_FAILED,
	SGX_FILE_STATUS_CRYPTO_ERROR,
	SGX_FILE_STATUS_CORRUPTED,
	SGX_FILE_STATUS_MEMORY_CORRUPTED,
	//SGX_FILE_STATUS_WRITE_TO_DISK_FAILED_NEED_MC,
	//SGX_FILE_STATUS_MC_NOT_INCREMENTED,
	SGX_FILE_STATUS_CLOSED,
} protected_fs_status_e;

#define MAX_PAGES_IN_CACHE 48

//COMPILE_TIME_ASSERT(filename_length, FILENAME_MAX_LEN == FILENAME_MAX);
COMPILE_TIME_ASSERT(filename_length, FILENAME_MAX_LEN <= FILENAME_MAX);

#ifdef NON_SGX_PROTECTED_FS
#else
typedef void FILE;
#endif

typedef union
{
	struct
	{
		uint8_t read   :1;
		uint8_t write  :1;
		uint8_t append :1;
		uint8_t binary: 1;
		uint8_t update: 1;
	};
	uint8_t raw;
} open_mode_t;


#define FILE_MHT_NODE_TYPE  1
#define FILE_DATA_NODE_TYPE 2

#define PATHNAME_MAX_LEN      (512)
#define FULLNAME_MAX_LEN      (PATHNAME_MAX_LEN+FILENAME_MAX_LEN)
#define RECOVERY_FILE_MAX_LEN (FULLNAME_MAX_LEN+10)

#pragma pack(push, 1)

/*
the following 2 structures are almost identical, i do not merge them (union or c++ inheritance from parent class) for 2 reasons:
1. when the code was written, they were much more different, during time the differences were almost gone, but "if it's working don't fix it", so i leave it this way
2. the code is more readable this way, it is clear when we deal with each type
*/
typedef struct _file_mht_node
{
	/* these are exactly the same as file_data_node_t below, any change should apply to both (both are saved in the cache as void*) */
	uint8_t type;
	uint64_t mht_node_number;
	struct _file_mht_node* parent;
	bool need_writing;
	bool new_node;
	union {
		struct {
			uint64_t physical_node_number;
			encrypted_node_t encrypted; // the actual data from the disk
		};
		recovery_node_t recovery_node;
	};
	/* from here the structures are different */
	mht_node_t plain; // decrypted data
} file_mht_node_t;


typedef struct _file_data_node
{
	/* these are exactly the same as file_mht_node_t above, any change should apply to both (both are saved in the cache as void*) */
	uint8_t type;
	uint64_t data_node_number;
	file_mht_node_t* parent;
	bool need_writing;
	bool new_node;
	union {
		struct {
			uint64_t physical_node_number;
			encrypted_node_t encrypted; // the actual data from the disk
		};
		recovery_node_t recovery_node;
	};
	/* from here the structures are different */
	data_node_t plain; // decrypted data
} file_data_node_t;


class protected_fs_file
{
private:
	union {
		struct {
			uint64_t meta_data_node_number; // for recovery purpose, so it is easy to write this node
			meta_data_node_t file_meta_data; // actual data from disk's meta data node
		};
		recovery_node_t meta_data_recovery_node;
	};

	meta_data_encrypted_t encrypted_part_plain; // encrypted part of meta data node, decrypted
	
	file_mht_node_t root_mht; // the root of the mht is always needed (for files bigger than 3KB)

	FILE* file; // OS's FILE pointer
	
	open_mode_t open_mode;
	uint8_t read_only;
	int64_t offset; // current file position (user's view)
	bool end_of_file; // flag

	int64_t real_file_size;
	
	bool need_writing; // flag
	uint32_t last_error; // last operation error
	protected_fs_status_e file_status;
	
#ifdef NON_SGX_PROTECTED_FS
	pthread_mutex_t mutex;
#else
	sgx_thread_mutex_t mutex;
#endif

	uint8_t use_user_kdk_key;
#ifdef NON_SGX_PROTECTED_FS
	aead_128bit_key_t user_kdk_key; // recieved from user, used instead of the seal key
	aead_128bit_key_t cur_key;
	aead_128bit_key_t session_master_key;
#else
	sgx_aes_gcm_128bit_key_t user_kdk_key; // recieved from user, used instead of the seal key
	sgx_aes_gcm_128bit_key_t cur_key;
	sgx_aes_gcm_128bit_key_t session_master_key;
#endif
	uint32_t master_key_count;
	
	char recovery_filename[RECOVERY_FILE_MAX_LEN]; // might include full path to the file

	lru_cache cache;

	// these don't change after init...
	sgx_iv_t empty_iv;

#ifdef NON_SGX_PROTECTED_FS
#else
	sgx_report_t report;
#endif

	void init_fields();
	bool cleanup_filename(const char* src, char* dest);
	bool parse_mode(const char* mode);
	bool file_recovery(const char* filename);
	bool init_existing_file(const char* filename, const char* clean_filename, const sgx_aes_gcm_128bit_key_t* import_key);
	bool init_new_file(const char* clean_filename);
	
	bool generate_secure_blob(sgx_aes_gcm_128bit_key_t* key, const char* label, uint64_t physical_node_number, sgx_aes_gcm_128bit_tag_t* output);
	bool generate_secure_blob_from_user_kdk(bool restore);
	bool init_session_master_key();
	bool derive_random_node_key(uint64_t physical_node_number);
	bool generate_random_meta_data_key();
	bool restore_current_meta_data_key(const sgx_aes_gcm_128bit_key_t* import_key);
	
	
	file_data_node_t* get_data_node();
	file_data_node_t* read_data_node();
	file_data_node_t* append_data_node();
	file_mht_node_t* get_mht_node();
	file_mht_node_t* read_mht_node(uint64_t mht_node_number);
	file_mht_node_t* append_mht_node(uint64_t mht_node_number);
	bool write_recovery_file();
	bool set_update_flag(bool flush_to_disk);
	void clear_update_flag();
	bool update_all_data_and_mht_nodes();
	bool update_meta_data_node();
	bool write_all_changes_to_disk(bool flush_to_disk);
	void erase_recovery_file();
	bool internal_flush(/*bool mc,*/ bool flush_to_disk);

public:
	protected_fs_file(const char* filename, const char* mode, const sgx_aes_gcm_128bit_key_t* import_key, const sgx_aes_gcm_128bit_key_t* kdk_key);
	~protected_fs_file();

	size_t write(const void* ptr, size_t size, size_t count);
	size_t read(void* ptr, size_t size, size_t count);
	int64_t tell();
	int seek(int64_t new_offset, int origin);
	bool get_eof();
	uint32_t get_error();
	void clear_error();
	int32_t clear_cache();
	bool flush(/*bool mc*/);
	bool pre_close(sgx_key_128bit_t* key, bool import);
	static int32_t remove(const char* filename);

	// Add for MesaTEE
	int32_t get_current_meta_gmac(sgx_aes_gcm_128bit_tag_t out_gmac);
	int32_t rename_meta(const char* old_name, const char* new_name);
};

#pragma pack(pop)

#endif // _PROTECTED_FS_H_
